# GA
import pre
import numpy as np
import numpy.random as rand
import random
import matplotlib.pyplot as plt
import time
import sys
import math as m
import xlrd
import itertools as itt

'''
consts for GA
'''
rand.seed()
g_n_grn = 0
# probabilities
n_popu = 300  # size of the population
crossover_prob = 0.85
mutation_prob = 0.15

n_cities = 12

len_chro = n_cities  # length of the chromosome. equals to #cities
max_grn = 300  # max generations(iterations)
g_dist = []
best_now = []   # the best individual of current generation, which should be gurantee to retain unless better one generated due to crossover or mutation


# the array of the probability every individual being selected according to the fitnesscurrent population(generation)
current_pop_accu_prob = []

'''
functions for generating
'''


def total_dist(perm):
    # the total distance of certain permulation
    total_dist = 0
    for i in range(n_cities-1):
        # i from 0 to n-1
        idx_a = perm[i]
        idx_b = perm[i + 1]
        total_dist = total_dist + g_dist[idx_a][idx_b]
    idx_a = perm[0]
    idx_b = perm[n_cities-1]
    total_dist = total_dist + g_dist[idx_a][idx_b]
    return (total_dist)


def fitness(perm):
    # the total distance of certain permulation
    total_dist = 0
    for i in range(n_cities-1):
        # i from 0 to n-1
        idx_a = perm[i]
        idx_b = perm[i + 1]
        total_dist = total_dist + g_dist[idx_a][idx_b]
    idx_a = perm[0]
    idx_b = perm[n_cities-1]
    total_dist = total_dist + g_dist[idx_a][idx_b]
    return (1/total_dist + 1e-6)


def associated_sorting(key, value):
    # https://www.geeksforgeeks.org/python-ways-to-sort-a-zipped-list-by-values/
    zipped = list(zip(key, value))
    # zip and unzip
    res = zip(*sorted(zipped, key=lambda x: x[1], reverse=True))
    key, value = res
    key, value = list(key), list(value)
    # print("KEY")
    # print(key)
    # print("VALUE")
    # print(value)
    return key, value


def popu_to_prob(popu, rank_based=False):
    # normalize fitness to prob
    # record the best individual of the current generation at the same time
    total = 0
    # sure to be replaced
    best_now = -1
    prob = []
    # print(prob)
    for i in range(n_popu):
        # for individuals in the population
        prob.append(fitness(popu[i]))
        # print(i)
        # print(prob)
        # record the best
        if prob[i] > best_now:
            best_now = prob[i]
        total = total + prob[i]

    accum = 0
    # associate the prob with the individuals
    popu, prob = associated_sorting(
        popu, prob)
    # higher, the better
    # print("IND")
    # print(popu)
    # print("PROB")
    # print(prob)

    print("CURRENT BEST of GRN " + str(g_n_grn))
    print("perm:", popu[0])
    print("fitness:", prob[0])

    current_pop_accu_prob = []

    if rank_based == False:
        for i in range(n_popu):
            # for individuals in the population
            # round - robin method
            accum = accum + prob[i] / total
            current_pop_accu_prob.append(accum)
    else:
        # rank_based
        nsum = n_popu*(n_popu+1)/2
        for i in range(n_popu-1, -1, -1):
            current_pop_accu_prob.append(i/nsum)

    # print("IND")
    # print(popu)

    # print("ACCU")
    # print(current_pop_accu_prob)
    # current_best = popu[0]

    current_best = popu[0]
    current_best_fit = prob[0]
    current_worst = popu[n_popu-1]
    current_worst_fit = prob[n_popu-1]
    current_best_sum = total_dist(popu[0])
    # current_pop_accu_prob

    return current_best, current_best_fit, current_best_sum, current_worst, current_worst_fit, current_pop_accu_prob, total


def pick_one(popu, current_pop_accu_prob):
    # pick one individual from the current generation(population)
    p = rand.uniform()
    for i in range(n_popu-1):
        if current_pop_accu_prob[i] <= p and p < current_pop_accu_prob[i + 1]:
            return popu[i]

    return popu[0]


def pick_better(ind1, ind2):
    # pick the better one from to
    fit1 = fitness(ind1)
    fit2 = fitness(ind2)
    return ind1 if fit1 > fit2 else ind2


def crossover(parent1, parent2):
    # crossover to generate offspring
    new_perm = pick_better(parent1, parent2)

    if random.random() < crossover_prob:
        # according to the probability
        new_perm = [0] * n_cities

        idxa = random.randint(0, n_cities-2)
        idxb = random.randint(idxa+1, n_cities-1)

        new_perm[idxa:idxb+1] = parent1[idxa:idxb+1]

        # copy
        remnant = parent2[:]

        for i in range(idxa, idxb+1):
            remnant.remove(parent1[i])
            # go through the unfilled blocks
        cur = 0
        for i in range(0, idxa):
            new_perm[i] = remnant[cur]
            cur = cur+1
        for i in range(idxb+1, n_cities):
            new_perm[i] = remnant[cur]
            cur = cur + 1
    return new_perm


def swap_mutate(ind):
    # individual is a permulation of n_cities
    # print("BEFORE MUTATION")
    # print(ind)

    if random.random() < mutation_prob:
        i = rand.randint(0, n_cities)
        j = (i + 1) % n_cities
        # swap in python
        ind[i], ind[j] = ind[j], ind[i]

    # print("AFTER MUTATION")
    # print(ind)


def insert_mutate(ind):
    # individual is a permulation of n_cities
    # print("BEFORE MUTATION")
    # print(ind)

    if random.random() < mutation_prob:
        i = rand.randint(0, n_cities)
        j = rand.randint(i+1, i+n_cities) % n_cities
        # insertion mutation
        temp = ind[i]
        ind.remove(temp)
        ind.insert(j, temp)
    # print("AFTER MUTATION")
    # print(ind)


def init_diagram(x, y, plt):

    x = x[0:n_cities]
    y = y[0:n_cities]

    fig = plt.figure()
    ax = fig.add_subplot(1, 1, 1)

    ax.scatter(x, y)
    ax.set_xlabel("Coord X")
    ax.set_ylabel("Coord Y")

    for i in range(n_cities):
        ax.annotate(str(i), (x[i], y[i]), color='black')

    std_text = ax.texts
    # print("TEXTS")
    # print(std_text)
    fig.canvas.draw()

    ax_bg = fig.canvas.copy_from_bbox(ax.bbox)

    plt.ion()
    return fig, ax, ax_bg, std_text


'''
functions for visualizing
'''


def plot(fig, ax, ax_bg, best_ever, best_sum, current_best, current_best_sum, std_text):

    # print(ind)
    # plot accroding to the permulation

    ax.lines = []
    # ax.texts = []
    
    fig.canvas.restore_region(ax_bg)

    # current_best
    xs = []
    ys = []
    for i in range(n_cities):
        ida = current_best[i]
        # idb = ind[i+1]
        xs.append(x[ida])
        ys.append(y[ida])

    xs.append(x[current_best[0]])
    ys.append(y[current_best[0]])
    p, = ax.plot(xs, ys, color='gray')
    ax.draw_artist(p)
    s = str(current_best)
    # best_ever
    xs = []
    ys = []
    for i in range(n_cities):
        ida = best_ever[i]
        # idb = ind[i+1]
        xs.append(x[ida])
        ys.append(y[ida])

    xs.append(x[best_ever[0]])
    ys.append(y[best_ever[0]])
    p, = ax.plot(xs, ys, color='r')
    ax.draw_artist(p)
    s = str(best_ever)

    title = ax.set_title(s+"\nSum of dist : "+str(best_sum), wrap=True)
    ax.draw_artist(title)



    fig.canvas.blit(ax.bbox)
# def plot(ind):

#     # plot accroding to the permulation

#     for i in range(n_cities-1):
#         ida = ind[i]
#         idb = ind[i+1]
#         plt.plot([x[ida], x[idb]], [y[ida], y[idb]], color='r')

    # for i in range(100000):
    #     y = np.random.random()
    #     x = np.random.random()
    #     plt.scatter(x, y)
    #     # plt.draw()
    #     plt.pause(1e-5)
    #     if not plt.get_fignums():
    #         sys.exit()


def ga(x, y):

    type_sel = 'prob'
    # type_sel = 'rank'
    type_mut = 'swap'
    # type_mut = 'isrt'
    names = '_typesel_'+type_sel+'_typemu_'+type_mut

    # s = 'E:/codes/courses/ArtificialIntelligence/GA/tsp_ga/out/out_cr_' + \
    #     str(crossover_prob) + '_mu_' + str(mutation_prob) + \
    #     '_npr_' + str(n_popu) + '_nct_' + str(n_cities) + names+'.txt'

    s = "E:/codes/courses/ArtificialIntelligence/GA/tsp_ga/out/test.txt"
    with open(s, 'w+') as file:
        sys.stdout = file  # 标准输出重定向至文件
        # sys.stdout = sys.stdout # 标准输出重定向至文件
        # x,y are the coords of cities
        # cut
        global g_n_grn

        fitness = []  # of n individuals
        best_ever = []
        worst_ever = []
        best_fitness = -1
        worst_fitness = 1e9
        total_fitness = 0
        popu = []  # current generation

        x = x[0:n_cities]
        y = y[0:n_cities]
        frame_count = 0

        fig, ax, ax_bg, std_text = init_diagram(x, y, plt)

        new_popu = []

        init_t = time.process_time()  # timing
        # generate the first population
        # randonmly init
        for i in range(n_popu):
            # for all of the individuals
            perm = list(rand.permutation(n_cities))
            new_popu.append(perm)

        # print(popu)
        # until reaches the max iteration

        for k in range(max_grn):
            g_n_grn = g_n_grn+1
            popu = new_popu
            new_popu = []
            current_pop_accu_prob = []
            current_best_fit = -1
            # loop until all next population generated completely
            current_best, current_best_fit, current_best_sum, current_worst, current_worst_fit, current_pop_accu_prob, current_total_fitness = popu_to_prob(
                popu)
            # current_best, current_best_fit, current_best_sum, current_worst, current_worst_fit, current_pop_accu_prob, current_total_fitness = popu_to_prob(popu,rank_based= True)

            # print(current_best)
            for i in range(n_popu):
                # the best two have the right to generate offspring
                p1 = pick_one(popu, current_pop_accu_prob)
                p2 = pick_one(popu, current_pop_accu_prob)

                offspring = crossover(p1, p2)
                swap_mutate(offspring)
                # insert_mutate(offspring)

                new_popu.append(offspring)
            if best_fitness < current_best_fit:
                best_fitness = current_best_fit
                best_ever = current_best
                best_sum = current_best_sum

            if current_worst_fit < worst_fitness:
                worst_fitness = current_worst_fit
                worst_ever = current_worst

            plot(fig, ax, ax_bg, best_ever, best_sum,
                 current_best, current_best_sum, std_text)
            total_fitness = total_fitness + current_total_fitness
            # plot(fig, ax, ax_bg, list(rand.permutation(n_cities)), best_sum, std_text)
            frame_count = frame_count + 1
            plt.pause(1e-3)

        print("**************GA DONE**************")
        print("TIME")
        print(time.process_time() - init_t)
        print("FRAME")
        print(frame_count)
        print("TOTAL FITNESS")
        print(total_fitness)
        print("AVG FITNESS")
        print(total_fitness/(max_grn*n_popu))
        print("WORST FITNESS")
        print(worst_fitness)
        print("BEST FITNESS")
        print(best_fitness)
    return best_ever, best_fitness


'''
functions to calculate the look-up table
'''


def cal_dist_matrix(file):

    # 经度
    x = []
    y = []
    dist = []

    wb = xlrd.open_workbook(filename=file)  # 打开文件
    print(wb.sheet_names())  # 获取所有表格名字

    sheet1 = wb.sheet_by_index(0)  # 通过索引获取表格
    print(sheet1.name, sheet1.nrows, sheet1.ncols)

    x = sheet1.col_values(0)
    y = sheet1.col_values(1)  # 获取列内容

    print(x)
    print(y)

    dist = np.zeros([len(x), len(x)], dtype=float)
    print(dist)

    # for i in range(10):
    lenx = len(x)
    for i in range(lenx):
        for j in range(lenx):
            try:
                dist[i, j] = m.sqrt((x[i]-x[j])**2+(y[i]-y[j])**2)
            except:
                print("************ERROR*************")
        print(i)

    np.savetxt(
        'E:/codes/courses/ArtificialIntelligence/GA/tsp_ga/dist.txt', dist, fmt='%0.8f')


# cal_dist_matrix("E:/codes/courses/ArtificialIntelligence/GA/raw eil101.xlsx")

def get_dist_matrix(file):
    dist = np.loadtxt(
        file, delimiter=' ')
    return dist


# dist = get_dist_matrix(
#     "E:/codes/courses/ArtificialIntelligence/GA/tsp_ga/dist.txt")
# # print(dist[0][1],dist[0][2])

'''
main function
'''

'''
testing codes
'''


# plt.ion()
g_dist = get_dist_matrix(
    "E:/codes/courses/ArtificialIntelligence/GA/tsp_ga/dist.txt")
xy = np.loadtxt(
    'E:/codes/courses/ArtificialIntelligence/GA/tsp_ga/raw_xy.txt')
x = xy[..., 0]
y = xy[..., 1]
# plt.ion()
# plt.show()
ga(x, y)
# x = [1, 4, 2, 8, 10,6]
# y = [[1, 2, 3], [2, 3, 1], [3, 1, 2], [1, 3, 2], [2, 1, 3], [3, 2, 1]]
# y, x = associated_sorting(y, x)
# print(y)
# print(x)

'''
testing code are above
'''
